#!/usr/bin/env python3
"""
创建测试用的RTSP流
使用OpenCV创建一个简单的测试视频流
"""

import cv2
import numpy as np
import time
import threading
import argparse
from datetime import datetime

class TestRTSPServer:
    """测试RTSP服务器（使用OpenCV模拟）"""
    
    def __init__(self, width=640, height=480, fps=30):
        self.width = width
        self.height = height
        self.fps = fps
        self.running = False
        self.frame_count = 0
        
    def generate_test_frame(self):
        """生成测试帧"""
        # 创建彩色背景
        frame = np.zeros((self.height, self.width, 3), dtype=np.uint8)
        
        # 渐变背景
        for y in range(self.height):
            for x in range(self.width):
                frame[y, x] = [
                    int(255 * x / self.width),  # R
                    int(255 * y / self.height), # G
                    128  # B
                ]
        
        # 添加移动的圆形（模拟检测对象）
        center_x = int(self.width/2 + 100 * np.sin(self.frame_count * 0.1))
        center_y = int(self.height/2 + 50 * np.cos(self.frame_count * 0.1))
        cv2.circle(frame, (center_x, center_y), 30, (255, 255, 255), -1)
        
        # 添加移动的矩形
        rect_x = int(100 + 200 * np.sin(self.frame_count * 0.05))
        rect_y = int(100 + 100 * np.cos(self.frame_count * 0.08))
        cv2.rectangle(frame, (rect_x, rect_y), (rect_x + 80, rect_y + 60), (0, 255, 255), -1)
        
        # 添加时间戳
        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        cv2.putText(frame, timestamp, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
        
        # 添加帧计数
        cv2.putText(frame, f"Frame: {self.frame_count}", (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
        
        self.frame_count += 1
        return frame
    
    def save_test_video(self, filename="test_video.mp4", duration=30):
        """保存测试视频文件"""
        print(f"🎬 创建测试视频: {filename}")
        print(f"   分辨率: {self.width}x{self.height}")
        print(f"   帧率: {self.fps} FPS")
        print(f"   时长: {duration} 秒")
        
        # 创建视频写入器
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        out = cv2.VideoWriter(filename, fourcc, self.fps, (self.width, self.height))
        
        total_frames = self.fps * duration
        
        for i in range(total_frames):
            frame = self.generate_test_frame()
            out.write(frame)
            
            if i % (self.fps * 5) == 0:  # 每5秒显示进度
                progress = (i / total_frames) * 100
                print(f"   进度: {progress:.1f}%")
        
        out.release()
        print(f"✅ 测试视频已保存: {filename}")
        return filename

def create_webcam_stream():
    """使用网络摄像头创建实时流"""
    print("📹 尝试打开网络摄像头...")
    
    # 尝试打开摄像头
    cap = cv2.VideoCapture(0)
    if not cap.isOpened():
        print("❌ 无法打开摄像头")
        return None
    
    print("✅ 摄像头已打开")
    print("按 'q' 键退出，按 's' 键截图")
    
    frame_count = 0
    
    while True:
        ret, frame = cap.read()
        if not ret:
            print("❌ 无法读取摄像头帧")
            break
        
        # 添加帧信息
        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        cv2.putText(frame, timestamp, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
        cv2.putText(frame, f"Frame: {frame_count}", (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
        
        # 显示帧
        cv2.imshow('Webcam Stream', frame)
        
        key = cv2.waitKey(1) & 0xFF
        if key == ord('q'):
            break
        elif key == ord('s'):
            screenshot_name = f"screenshot_{frame_count}.jpg"
            cv2.imwrite(screenshot_name, frame)
            print(f"📸 截图已保存: {screenshot_name}")
        
        frame_count += 1
    
    cap.release()
    cv2.destroyAllWindows()
    print("📹 摄像头流已关闭")

def play_video_file(filename):
    """播放视频文件"""
    print(f"🎬 播放视频文件: {filename}")
    
    cap = cv2.VideoCapture(filename)
    if not cap.isOpened():
        print(f"❌ 无法打开视频文件: {filename}")
        return
    
    fps = cap.get(cv2.CAP_PROP_FPS)
    frame_delay = 1.0 / fps if fps > 0 else 1.0 / 30
    
    print(f"   视频FPS: {fps}")
    print("按 'q' 键退出，按空格键暂停/继续")
    
    paused = False
    
    while True:
        if not paused:
            ret, frame = cap.read()
            if not ret:
                print("🔄 视频播放完毕，重新开始")
                cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
                continue
            
            cv2.imshow('Video Player', frame)
        
        key = cv2.waitKey(int(frame_delay * 1000)) & 0xFF
        if key == ord('q'):
            break
        elif key == ord(' '):
            paused = not paused
            print(f"⏸️ 视频{'暂停' if paused else '继续'}")
    
    cap.release()
    cv2.destroyAllWindows()
    print("🎬 视频播放已停止")

def main():
    parser = argparse.ArgumentParser(description="创建测试用的视频流")
    parser.add_argument("--mode", choices=["video", "webcam", "play"], default="video",
                       help="模式: video=创建测试视频, webcam=使用摄像头, play=播放视频")
    parser.add_argument("--output", default="test_video.mp4", help="输出视频文件名")
    parser.add_argument("--input", help="输入视频文件名（play模式）")
    parser.add_argument("--width", type=int, default=640, help="视频宽度")
    parser.add_argument("--height", type=int, default=480, help="视频高度")
    parser.add_argument("--fps", type=int, default=30, help="帧率")
    parser.add_argument("--duration", type=int, default=30, help="视频时长（秒）")
    
    args = parser.parse_args()
    
    if args.mode == "video":
        # 创建测试视频
        server = TestRTSPServer(args.width, args.height, args.fps)
        server.save_test_video(args.output, args.duration)
        
        print(f"\n📝 使用说明:")
        print(f"1. 测试视频已创建: {args.output}")
        print(f"2. 可以使用以下命令播放:")
        print(f"   python create_test_rtsp.py --mode play --input {args.output}")
        print(f"3. 或者使用RTSP测试客户端:")
        print(f"   python test_rtsp_client.py --rtsp {args.output}")
        
    elif args.mode == "webcam":
        # 使用摄像头
        create_webcam_stream()
        
    elif args.mode == "play":
        # 播放视频
        if not args.input:
            print("❌ play模式需要指定 --input 参数")
            return
        play_video_file(args.input)
    
    print("✅ 完成")

if __name__ == "__main__":
    main()
